import numpy as np

class HMMAdapter:
    def __init__(self, model, obs_ops, tseq, x_b, r_matrix):
        self.Dyn = ModelAdapter(model)
        self._obs = [ObsAdapter(op, r_matrix) for op in obs_ops]
        self.Obs = lambda ko: self._obs[ko] if 0 <= ko <= len(self._obs)-1 else None
        self.tseq = TimeSeqAdapter(tseq)
        self.X0 = InitialStateAdapter(x_b)
        
class ModelAdapter:
    def __init__(self, model):
        self.model = model
        self.M = model.state_size
        
    def __call__(self, x, t, dt):
        return self.model.step(x)
        
    def linear(self, x, t, dt):
        # Validate input dimensions
        if x.size != self.M:
            raise ValueError(f"State vector size {x.size} doesn't match model state size {self.M}")
            
        # Handle 2D state for barotropic model
        if hasattr(self.model, 'shape'):
            x_2d = x.reshape(self.model.shape)
            M = self.model.tangent_linear(x_2d)
        else:
            M = self.model.tangent_linear(x)
            
        # Validate tangent linear matrix dimensions
        if M.shape != (self.M, self.M):
            raise ValueError(f"Tangent linear matrix shape {M.shape} doesn't match expected ({self.M}, {self.M})")
            
        return M

class ObsAdapter:
    def __init__(self, obs_op, r_matrix):
        self.obs_op = obs_op
        self.noise = NoiseAdapter(r_matrix)
        
    def __call__(self, x):
        """Apply observation operator or return self based on input type"""
        if isinstance(x, (int, np.integer)) or x is None:
            return self
        else:
            # Handle 2D state for barotropic model
            if hasattr(self.obs_op, 'model') and hasattr(self.obs_op.model, 'shape'):
                x_2d = x.reshape(self.obs_op.model.shape)
                return self.obs_op.h(x_2d)
            return self.obs_op.h(x)
        
    def linear(self, x):
        """Return H matrix explicitly as 2D array"""
        # Handle 2D state for barotropic model
        if hasattr(self.obs_op, 'model') and hasattr(self.obs_op.model, 'shape'):
            x_2d = x.reshape(self.obs_op.model.shape)
            H = self.obs_op.h_linearized(x_2d)
        else:
            H = self.obs_op.h_linearized(x)
            
        # Ensure proper 2D shape with correct dimensions
        if H.ndim < 2:
            H = np.atleast_2d(H)
        if H.shape[0] != self.obs_op.obs_size:
            H = H.reshape(self.obs_op.obs_size, -1)
        return H

class NoiseAdapter:
    def __init__(self, cov):
        # Ensure diagonal covariance matrix format
        if cov.ndim == 1:
            self.C = CovMatrixAdapter(np.diag(cov))
        else:
            self.C = CovMatrixAdapter(cov)

class CovMatrixAdapter:
    def __init__(self, cov):
        self.matrix = np.atleast_2d(cov)
        
    @property
    def sym_sqrt_inv(self):
        """Symmetric square root inverse using eigenvalue decomposition"""
        s, U = np.linalg.eigh(self.matrix)
        # Handle positive semi-definite
        sqrt_s = np.sqrt(np.where(s > 0, s, 0))
        inv_sqrt_s = 1.0 / np.where(sqrt_s > 0, sqrt_s, 1)
        return U @ np.diag(inv_sqrt_s) @ U.T

class TimeSeqAdapter:
    def __init__(self, time_window):
        self.Ko = time_window
        self.kko = {k: k for k in range(time_window+1)}
        
    def cycle(self, ko):
        # Only yield valid time steps
        if 0 <= ko <= self.Ko:
            # Yield (k, t, dt) for a single cycle
            yield ko, ko * 0.1, 0.1

class InitialStateAdapter:
    def __init__(self, x0):
        self.mu = x0
